#include <iostream>
#include <vector>
#include <cmath>
#include <Eigen/Dense>
#include <limits>
class FormationMatcher {
public:
    FormationMatcher(int num)
        : n(num),
          adj_matrix(n, n),
          label_left(n, 0),
          label_right(n, 0),
          match_right(n, -1),
          slack_right(n, std::numeric_limits<double>::max()),
          visit_left(n, false),
          visit_right(n, false) 
    {
        std::cout << "formation matcher" << std::endl;
    }

    std::vector<int> match() {
        // 遍历每个左顶点（每架无人机）
        for (int i = 0; i < n; ++i) {
            // 初始化右顶点的松弛量数组
            std::fill(slack_right.begin(), slack_right.end(), std::numeric_limits<double>::max());

            // 不断尝试寻找增广路径，直到找到为止
            while (true) {
                // 重置访问标记数组
                std::fill(visit_left.begin(), visit_left.end(), false);
                std::fill(visit_right.begin(), visit_right.end(), false);

                // 尝试寻找增广路径
                if (find_path(i)) {
                    break;
                }

                // 计算最小松弛量 d
                double d = std::numeric_limits<double>::infinity();
                for (int j = 0; j < n; ++j) {
                    if (!visit_right[j]) {
                        d = std::min(d, slack_right[j]);
                    }
                }

                // 更新标签和松弛量
                for (int k = 0; k < n; ++k) {
                    if (visit_left[k]) label_left[k] -= d;
                    if (visit_right[k]) label_right[k] += d;
                    else slack_right[k] -= d;
                }
            }
        }
        std::cout << "match finished!" << std::endl;
        // 返回匹配结果
        return match_right;
    }

    void set_values(const std::vector<Eigen::Vector3d> initial_formations, const std::vector<Eigen::Vector3d>target_formations){
        initial_formation.clear();
        target_formation.clear();
        for (int i = 0; i < n; ++i) {
            initial_formation.push_back(initial_formations[i]);
            target_formation.push_back(target_formations[i]);
        }
    }
    void init(){
        for (int i = 0; i < n; ++i) {
            for (int j = 0; j < n; ++j) {
                adj_matrix(i, j) = -1.0  *(initial_formation[i] - target_formation[j]).norm();
            }
        }

        // 初始化左顶点标签为每行的最大值
        for (int i = 0; i < n; ++i) {
            label_left[i] = adj_matrix.row(i).maxCoeff();
        }
    }
private:
    int n; // 无人机数量或目标位置数量
    Eigen::MatrixXd adj_matrix; // 邻接矩阵，存储无人机到目标位置的距离
    std::vector<double> label_left; // 左顶点标签（无人机标签）
    std::vector<double> label_right; // 右顶点标签（目标位置标签）
    std::vector<int> match_right; // 每个目标位置的匹配结果，存储匹配的无人机编号
    std::vector<double> slack_right; // 右顶点的松弛量
    std::vector<bool> visit_left; // 左顶点访问标记
    std::vector<bool> visit_right; // 右顶点访问标记
    std::vector<Eigen::Vector3d>initial_formation;
    std::vector<Eigen::Vector3d>target_formation;

    bool find_path(int i) {
        visit_left[i] = true;

        for (int j = 0; j < n; ++j) {
            if (visit_right[j]) continue;

            double gap = label_left[i] + label_right[j] - adj_matrix(i, j);

            if (gap == 0) {
                visit_right[j] = true;
                if (match_right[j] == -1 || find_path(match_right[j])) {
                    match_right[j] = i;
                    return true;
                }
            } else {
                slack_right[j] = std::min(slack_right[j], gap);
            }
        }
        return false;
    }
};